import numpy as np
from numpy import nan
import pandas as pd

import pytest
from numpy.testing import assert_allclose

import pvlib
from pvlib import tracking
from .conftest import TESTS_DATA_DIR, assert_frame_equal, assert_series_equal

SINGLEAXIS_COL_ORDER = ['tracker_theta', 'aoi',
                        'surface_azimuth', 'surface_tilt']


def test_solar_noon():
    index = pd.date_range(start='20180701T1200', freq='1s', periods=1)
    apparent_zenith = pd.Series([10], index=index)
    apparent_azimuth = pd.Series([180], index=index)
    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'tracker_theta': 0, 'aoi': 10,
                           'surface_azimuth': 90, 'surface_tilt': 0},
                          index=index, dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)


def test_scalars():
    apparent_zenith = 10
    apparent_azimuth = 180
    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)
    assert isinstance(tracker_data, dict)
    expect = {'tracker_theta': 0, 'aoi': 10, 'surface_azimuth': 90,
              'surface_tilt': 0}
    for k, v in expect.items():
        assert np.isclose(tracker_data[k], v)


def test_arrays():
    apparent_zenith = np.array([10])
    apparent_azimuth = np.array([180])
    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)
    assert isinstance(tracker_data, dict)
    expect = {'tracker_theta': 0, 'aoi': 10, 'surface_azimuth': 90,
              'surface_tilt': 0}
    for k, v in expect.items():
        assert_allclose(tracker_data[k], v, atol=1e-7)


def test_nans():
    apparent_zenith = np.array([10, np.nan, 10])
    apparent_azimuth = np.array([180, 180, np.nan])
    with np.errstate(invalid='ignore'):
        tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                           axis_tilt=0, axis_azimuth=0,
                                           max_angle=90, backtrack=True,
                                           gcr=2.0/7.0)
    expect = {'tracker_theta': np.array([0, nan, nan]),
              'aoi': np.array([10, nan, nan]),
              'surface_azimuth': np.array([90, nan, nan]),
              'surface_tilt': np.array([0, nan, nan])}
    for k, v in expect.items():
        assert_allclose(tracker_data[k], v, atol=1e-7)

    # repeat with Series because nans can differ
    apparent_zenith = pd.Series(apparent_zenith)
    apparent_azimuth = pd.Series(apparent_azimuth)
    with np.errstate(invalid='ignore'):
        tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                           axis_tilt=0, axis_azimuth=0,
                                           max_angle=90, backtrack=True,
                                           gcr=2.0/7.0)
    expect = pd.DataFrame(np.array(
        [[ 0., 10., 90.,  0.],
         [nan, nan, nan, nan],
         [nan, nan, nan, nan]]),
        columns=['tracker_theta', 'aoi', 'surface_azimuth', 'surface_tilt'])
    assert_frame_equal(tracker_data, expect)


def test_arrays_multi():
    apparent_zenith = np.array([[10, 10], [10, 10]])
    apparent_azimuth = np.array([[180, 180], [180, 180]])
    # singleaxis should fail for num dim > 1
    with pytest.raises(ValueError):
        tracking.singleaxis(apparent_zenith, apparent_azimuth,
                            axis_tilt=0, axis_azimuth=0,
                            max_angle=90, backtrack=True,
                            gcr=2.0/7.0)
    # uncomment if we ever get singleaxis to support num dim > 1 arrays
    # assert isinstance(tracker_data, dict)
    # expect = {'tracker_theta': np.full_like(apparent_zenith, 0),
    #           'aoi': np.full_like(apparent_zenith, 10),
    #           'surface_azimuth': np.full_like(apparent_zenith, 90),
    #           'surface_tilt': np.full_like(apparent_zenith, 0)}
    # for k, v in expect.items():
    #     assert_allclose(tracker_data[k], v)


def test_azimuth_north_south():
    apparent_zenith = pd.Series([60])
    apparent_azimuth = pd.Series([90])

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=180,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'tracker_theta': -60, 'aoi': 0,
                           'surface_azimuth': 90, 'surface_tilt': 60},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect['tracker_theta'] *= -1

    assert_frame_equal(expect, tracker_data)


def test_max_angle():
    apparent_zenith = pd.Series([60])
    apparent_azimuth = pd.Series([90])
    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=45, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 15, 'surface_azimuth': 90,
                           'surface_tilt': 45, 'tracker_theta': 45},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)


def test_min_angle():
    apparent_zenith = pd.Series([60])
    apparent_azimuth = pd.Series([270])
    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=(-45, 50), backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 15, 'surface_azimuth': 270,
                           'surface_tilt': 45, 'tracker_theta': -45},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)


def test_backtrack():
    apparent_zenith = pd.Series([80])
    apparent_azimuth = pd.Series([90])

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=90, backtrack=False,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 0, 'surface_azimuth': 90,
                           'surface_tilt': 80, 'tracker_theta': 80},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=0,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 52.5716, 'surface_azimuth': 90,
                           'surface_tilt': 27.42833, 'tracker_theta': 27.4283},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)


def test_axis_tilt():
    apparent_zenith = pd.Series([30])
    apparent_azimuth = pd.Series([135])

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=30, axis_azimuth=180,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 7.286245, 'surface_azimuth': 142.65730,
                           'surface_tilt': 35.98741,
                           'tracker_theta': -20.88121},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=30, axis_azimuth=0,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 47.6632, 'surface_azimuth': 50.96969,
                           'surface_tilt': 42.5152, 'tracker_theta': 31.6655},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)


def test_axis_azimuth():
    apparent_zenith = pd.Series([30])
    apparent_azimuth = pd.Series([90])

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=90,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 30, 'surface_azimuth': 180,
                           'surface_tilt': 0, 'tracker_theta': 0},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)

    apparent_zenith = pd.Series([30])
    apparent_azimuth = pd.Series([180])

    tracker_data = tracking.singleaxis(apparent_zenith, apparent_azimuth,
                                       axis_tilt=0, axis_azimuth=90,
                                       max_angle=90, backtrack=True,
                                       gcr=2.0/7.0)

    expect = pd.DataFrame({'aoi': 0, 'surface_azimuth': 180,
                           'surface_tilt': 30, 'tracker_theta': 30},
                          index=[0], dtype=np.float64)
    expect = expect[SINGLEAXIS_COL_ORDER]

    assert_frame_equal(expect, tracker_data)


def test_horizon_flat():
    # GH 569
    solar_azimuth = np.array([0, 180, 359])
    solar_zenith = np.array([100, 45, 100])
    solar_azimuth = pd.Series(solar_azimuth)
    solar_zenith = pd.Series(solar_zenith)
    # depending on platform and numpy versions this will generate
    # RuntimeWarning: invalid value encountered in > < >=
    out = tracking.singleaxis(solar_zenith, solar_azimuth, axis_tilt=0,
                              axis_azimuth=180, backtrack=False, max_angle=180)
    expected = pd.DataFrame(np.array(
        [[ nan,  nan,  nan,  nan],
         [  0.,  45., 270.,   0.],
         [ nan,  nan,  nan,  nan]]),
        columns=['tracker_theta', 'aoi', 'surface_azimuth', 'surface_tilt'])
    assert_frame_equal(out, expected)


def test_horizon_tilted():
    # GH 569
    solar_azimuth = np.array([0, 180, 359])
    solar_zenith = np.full_like(solar_azimuth, 45)
    solar_azimuth = pd.Series(solar_azimuth)
    solar_zenith = pd.Series(solar_zenith)
    out = tracking.singleaxis(solar_zenith, solar_azimuth, axis_tilt=90,
                              axis_azimuth=180, backtrack=False, max_angle=180)
    expected = pd.DataFrame(np.array(
        [[-180.,  45.,   0.,  90.],
         [   0.,  45., 180.,  90.],
         [ 179.,  45., 359.,  90.]]),
        columns=['tracker_theta', 'aoi', 'surface_azimuth', 'surface_tilt'])
    assert_frame_equal(out, expected)


def test_low_sun_angles():
    # GH 656, 824
    result = tracking.singleaxis(
        apparent_zenith=80, solar_azimuth=338, axis_tilt=30,
        axis_azimuth=180, max_angle=60, backtrack=True, gcr=0.35)
    expected = {
        'tracker_theta': np.array([60.0]),
        'aoi': np.array([80.420987]),
        'surface_azimuth': np.array([253.897886]),
        'surface_tilt': np.array([64.341094])}
    for k, v in result.items():
        assert_allclose(expected[k], v)


def test_calc_axis_tilt():
    # expected values
    expected_axis_tilt = 2.239  # [degrees]
    expected_side_slope = 9.86649274360294  # [degrees]
    expected = TESTS_DATA_DIR / 'singleaxis_tracker_wslope.csv'
    expected = pd.read_csv(expected, index_col='timestamp', parse_dates=True)
    # solar positions
    starttime = '2017-01-01T00:30:00-0300'
    stoptime = '2017-12-31T23:59:59-0300'
    lat, lon = -27.597300, -48.549610
    times = pd.DatetimeIndex(pd.date_range(starttime, stoptime, freq='h'))
    solpos = pvlib.solarposition.get_solarposition(times, lat, lon)
    # singleaxis tracker w/slope data
    slope_azimuth, slope_tilt = 77.34, 10.1149
    axis_azimuth = 0.0
    max_angle = 75.0
    # Note: GCR is relative to horizontal distance between rows
    gcr = 0.33292759  # GCR = length / horizontal_pitch = 1.64 / 5 / cos(9.86)
    # calculate tracker axis zenith
    axis_tilt = tracking.calc_axis_tilt(
        slope_azimuth, slope_tilt, axis_azimuth=axis_azimuth)
    assert np.isclose(axis_tilt, expected_axis_tilt)
    # calculate cross-axis tilt and relative rotation
    cross_axis_tilt = tracking.calc_cross_axis_tilt(
        slope_azimuth, slope_tilt, axis_azimuth, axis_tilt)
    assert np.isclose(cross_axis_tilt, expected_side_slope)
    sat = tracking.singleaxis(
        solpos.apparent_zenith, solpos.azimuth, axis_tilt, axis_azimuth,
        max_angle, backtrack=True, gcr=gcr, cross_axis_tilt=cross_axis_tilt)
    np.testing.assert_allclose(
        sat['tracker_theta'], expected['tracker_theta'], atol=1e-7)
    np.testing.assert_allclose(sat['aoi'], expected['aoi'], atol=1e-7)
    np.testing.assert_allclose(
        sat['surface_azimuth'], expected['surface_azimuth'], atol=1e-7)
    np.testing.assert_allclose(
        sat['surface_tilt'], expected['surface_tilt'], atol=1e-7)


def test_slope_aware_backtracking():
    """
    Test validation data set from https://www.nrel.gov/docs/fy20osti/76626.pdf
    """
    index = pd.date_range('2019-01-01T08:00', '2019-01-01T17:00', freq='h')
    index = index.tz_localize('Etc/GMT+5')
    expected_data = pd.DataFrame(index=index, data=[
        ( 2.404287, 122.79177, -84.440, -10.899),
        (11.263058, 133.288729, -72.604, -25.747),
        (18.733558, 145.285552, -59.861, -59.861),
        (24.109076, 158.939435, -45.578, -45.578),
        (26.810735, 173.931802, -28.764, -28.764),
        (26.482495, 189.371536, -8.475, -8.475),
        (23.170447, 204.13681, 15.120, 15.120),
        (17.296785, 217.446538, 39.562, 39.562),
        ( 9.461862, 229.102218, 61.587, 32.339),
        ( 0.524817, 239.330401, 79.530, 5.490),
    ], columns=['ApparentElevation', 'SolarAzimuth',
                'TrueTracking', 'Backtracking'])
    expected_axis_tilt = 9.666
    expected_slope_angle = -2.576
    slope_azimuth, slope_tilt = 180.0, 10.0
    axis_azimuth = 195.0
    axis_tilt = tracking.calc_axis_tilt(
        slope_azimuth, slope_tilt, axis_azimuth)
    assert np.isclose(axis_tilt, expected_axis_tilt, rtol=1e-3, atol=1e-3)
    cross_axis_tilt = tracking.calc_cross_axis_tilt(
        slope_azimuth, slope_tilt, axis_azimuth, axis_tilt)
    assert np.isclose(
        cross_axis_tilt, expected_slope_angle, rtol=1e-3, atol=1e-3)
    sat = tracking.singleaxis(
        90.0-expected_data['ApparentElevation'], expected_data['SolarAzimuth'],
        axis_tilt, axis_azimuth, max_angle=90.0, backtrack=True, gcr=0.5,
        cross_axis_tilt=cross_axis_tilt)
    assert_series_equal(sat['tracker_theta'],
                        expected_data['Backtracking'].rename('tracker_theta'),
                        check_less_precise=True)
    truetracking = tracking.singleaxis(
        90.0-expected_data['ApparentElevation'], expected_data['SolarAzimuth'],
        axis_tilt, axis_azimuth, max_angle=90.0, backtrack=False, gcr=0.5,
        cross_axis_tilt=cross_axis_tilt)
    assert_series_equal(truetracking['tracker_theta'],
                        expected_data['TrueTracking'].rename('tracker_theta'),
                        check_less_precise=True)


def test_singleaxis_aoi_gh1221():
    # vertical tracker
    loc = pvlib.location.Location(40.1134, -88.3695)
    dr = pd.date_range(
        start='02-Jun-1998 00:00:00', end='02-Jun-1998 23:55:00', freq='5min',
        tz='Etc/GMT+6')
    sp = loc.get_solarposition(dr)
    tr = pvlib.tracking.singleaxis(
        sp['apparent_zenith'], sp['azimuth'], axis_tilt=90, axis_azimuth=180,
        max_angle=0.001, backtrack=False)
    fixed = pvlib.irradiance.aoi(90, 180, sp['apparent_zenith'], sp['azimuth'])
    fixed[np.isnan(tr['aoi'])] = np.nan
    assert np.allclose(tr['aoi'], fixed, equal_nan=True)


def test_calc_surface_orientation_types():
    # numpy arrays
    rotations = np.array([-10, 0, 10])
    expected_tilts = np.array([10, 0, 10], dtype=float)
    expected_azimuths = np.array([270, 90, 90], dtype=float)
    out = tracking.calc_surface_orientation(tracker_theta=rotations)
    np.testing.assert_allclose(expected_tilts, out['surface_tilt'])
    np.testing.assert_allclose(expected_azimuths, out['surface_azimuth'])

    # pandas Series
    rotations = pd.Series(rotations)
    expected_tilts = pd.Series(expected_tilts).rename('surface_tilt')
    expected_azimuths = pd.Series(expected_azimuths).rename('surface_azimuth')
    out = tracking.calc_surface_orientation(tracker_theta=rotations)
    assert_series_equal(expected_tilts, out['surface_tilt'])
    assert_series_equal(expected_azimuths, out['surface_azimuth'])

    # float
    for rotation, expected_tilt, expected_azimuth in zip(
            rotations, expected_tilts, expected_azimuths):
        out = tracking.calc_surface_orientation(rotation)
        assert out['surface_tilt'] == pytest.approx(expected_tilt)
        assert out['surface_azimuth'] == pytest.approx(expected_azimuth)


def test_calc_surface_orientation_kwargs():
    # non-default axis tilt & azimuth
    rotations = np.array([-10, 0, 10])
    expected_tilts = np.array([22.2687445, 20.0, 22.2687445])
    expected_azimuths = np.array([152.72683041, 180.0, 207.27316959])
    out = tracking.calc_surface_orientation(rotations,
                                            axis_tilt=20,
                                            axis_azimuth=180)
    np.testing.assert_allclose(out['surface_tilt'], expected_tilts)
    np.testing.assert_allclose(out['surface_azimuth'], expected_azimuths)


def test_calc_surface_orientation_special():
    # special cases for rotations
    rotations = np.array([-180, -90, -0, 0, 90, 180])
    expected_tilts = np.array([180, 90, 0, 0, 90, 180], dtype=float)
    expected_azimuths = [270, 270, 90, 90, 90, 90]
    out = tracking.calc_surface_orientation(rotations)
    np.testing.assert_allclose(out['surface_tilt'], expected_tilts)
    np.testing.assert_allclose(out['surface_azimuth'], expected_azimuths)

    # special case for axis_tilt
    rotations = np.array([-10, 0, 10])
    expected_tilts = np.array([90, 90, 90], dtype=float)
    expected_azimuths = np.array([350, 0, 10], dtype=float)
    out = tracking.calc_surface_orientation(rotations, axis_tilt=90)
    np.testing.assert_allclose(out['surface_tilt'], expected_tilts)
    np.testing.assert_allclose(out['surface_azimuth'], expected_azimuths)

    # special cases for axis_azimuth
    rotations = np.array([-10, 0, 10])
    expected_tilts = np.array([10, 0, 10], dtype=float)
    expected_azimuth_offsets = np.array([-90, 90, 90], dtype=float)
    for axis_azimuth in [0, 90, 180, 270, 360]:
        expected_azimuths = (axis_azimuth + expected_azimuth_offsets) % 360
        out = tracking.calc_surface_orientation(rotations,
                                                axis_azimuth=axis_azimuth)
        np.testing.assert_allclose(out['surface_tilt'], expected_tilts)
        # the rounding is a bit ugly, but necessary to test approximately equal
        # in a modulo-360 sense.
        np.testing.assert_allclose(np.round(out['surface_azimuth'], 4) % 360,
                                   expected_azimuths, rtol=1e-5, atol=1e-5)
